package com.example.newdemo05;

import android.content.Context;
import android.content.Intent;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.os.Bundle;
import android.util.Log;
import android.view.View;
import android.widget.Button;
import android.widget.ImageView;
import android.widget.TextView;

import org.pytorch.IValue;
import org.pytorch.Module;
import org.pytorch.Tensor;
import org.pytorch.torchvision.TensorImageUtils;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

import androidx.appcompat.app.AppCompatActivity;

import com.example.newdemo05.classify.ImageNetClasses;

public class MainActivity extends AppCompatActivity {
    Button takePictureBtn=null;
    Bitmap bitmap = null;
    Module module = null;
    // 将图片显示在界面上
    ImageView imageView = null;
    TextView textView = null;
    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);
        initView();
        //从assets中获得图片数据
        try {
            bitmap = BitmapFactory.decodeStream(getAssets().open("image.jpg"));
            imgClassify(bitmap);
        } catch (IOException e) {
            Log.e("PytorchHelloWorld", "Error reading assets", e);
            finish();
        }


    }

    private void imgClassify(Bitmap img){
        try {
            // 加载PyTorch序列化模型
            module = Module.load(assetFilePath(this, "model101.pt"));
        } catch (IOException e) {
            Log.e("PytorchHelloWorld", "Error reading assets", e);
            finish();
        }
        imageView.setImageBitmap(img);
        // 建立输入张量
        final Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(img,
                TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB);
        // 运行模型推理
        final Tensor outputTensor = module.forward(IValue.from(inputTensor)).toTensor();
        // 获取推理结果
        final float[] scores = outputTensor.getDataAsFloatArray();
        //  获取概率最高的分类的索引号
        float maxScore = -Float.MAX_VALUE;
        int maxScoreIdx = -1;
        for (int i = 0; i < scores.length; i++) {
            if (scores[i] > maxScore) {
                maxScore = scores[i];
                maxScoreIdx = i;
            }
        }
        //通过索引号获得分类名称
        String className = ImageNetClasses.IMAGENET_CLASSES[maxScoreIdx];
        textView.setText("识别结果为:"+className);
    }

    private void initView(){
        takePictureBtn = findViewById(R.id.button);
        // 将图片显示在界面上
        imageView = findViewById(R.id.image);
        // 显示分类名称
        textView = findViewById(R.id.text);

        takePictureBtn.setOnClickListener(new View.OnClickListener() {
            @Override
            public void onClick(View v) {
                Intent intent = new Intent(MainActivity.this,TakePicturesActivity.class);
                startActivity(intent);
            }
        });
    }

    //返回模型文件路径
    public static String assetFilePath(Context context, String assetName) throws IOException {
        File file = new File(context.getFilesDir(), assetName);
        if (file.exists() && file.length() > 0) {
            return file.getAbsolutePath();
        }

        try (InputStream is = context.getAssets().open(assetName)) {
            try (OutputStream os = new FileOutputStream(file)) {
                byte[] buffer = new byte[4 * 1024];
                int read;
                while ((read = is.read(buffer)) != -1) {
                    os.write(buffer, 0, read);
                }
                os.flush();
            }
            return file.getAbsolutePath();
        }
    }
}
